import importlib
import numpy as np
import pandas as pd
from configs.config_loader import Config  # 假设你已有这个模块
import envs.future_env as future_env  # 你的环境模块
import os
import presentation.show as show
import utils.utils as utils
import sys, os
import torch

project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
if project_root not in sys.path:
    sys.path.insert(0, project_root)


def backtesting(data: np.ndarray, strategy: str):
    """
    Args:
        data: np.ndarray of shape (b, l, 20)
        strategy: str, e.g. "MACD_RSI_OBV_Strategy"
    """
    b, l, c = data.shape
    assert c == 20, "Expected 20 features per timestep"

    # === 动态加载策略模块 ===
    module = importlib.import_module(f"strategies.{strategy}")
    predict_action = getattr(module, "predict_action")

    # === 提取并保存 price_file 和 tech_file ===
    close = data[:, :, 0]  # shape: (b, l)
    tech = data[:, :, 1:]  # shape: (b, l, 19)

    # 转置为 (l, b) 格式保存为 price
    price_df = pd.DataFrame(close.T)  # shape: (l, b)
    tech_reshaped = tech.transpose(1, 0, 2).reshape(l, b * 19)  # shape: (l, b*19)
    tech_df = pd.DataFrame(tech_reshaped)

    # === 加载配置并构造环境 ===
    cfg = Config(yaml_path="baseline/configs/config_backtesting.yaml")
    env_config = cfg.env.copy()  # 假设 cfg.env 是字典
    env_config["price_file_path"] = price_df
    env_config["tech_file_path"] = tech_df

    env = future_env.MyFuturesEnv(env_config)
    env.set_test()
    state = env.reset()
    done = False
    transition_dict = {
        "states": [],
        "actions": [],
        "next_states": [],
        "rewards": [],
        "dones": [],
        "portfolio": [],
    }
    while not done:
        state = utils.process_state_to_blc(
            state,
            env.get_window(),
            env.get_future_count(),
            env.get_tech_num(),
        )
        action = predict_action(state)  # shape: (b, 1)
        if isinstance(action, np.ndarray):
            action = torch.from_numpy(action)
        next_state, reward, done, _ = env.step(action)
        transition_dict["states"].append(state)
        transition_dict["actions"].append(action)
        transition_dict["dones"].append(done)
        transition_dict["portfolio"].append(env.get_final_ratio())
        state = next_state

    return (
        transition_dict["portfolio"],
        env.get_price_array(),
        transition_dict["actions"],
    )


def clean_data(data):
    """清洗数据，确保所有元素是 float 或 int"""
    cleaned = []
    for item in data:
        if isinstance(item, (list, np.ndarray)):  # 如果是列表/数组，取第一个元素
            cleaned.append(float(item[0]))
        elif isinstance(item, (float, int)):  # 如果是标量，直接保留
            cleaned.append(float(item))
        else:  # 其他情况（如字符串），尝试转换
            try:
                cleaned.append(float(item))
            except ValueError:
                cleaned.append(np.nan)  # 无效数据填充 NaN
    return np.array(cleaned, dtype=np.float32)


if __name__ == "__main__":
    price_file_path = "data/output/A_tech_price.csv"
    tech_file_path = "data/output/A_tech_tech.csv"
    price_array = pd.read_csv(price_file_path).to_numpy()
    tech_array = pd.read_csv(tech_file_path).to_numpy()

    data = np.hstack((price_array, tech_array))
    data = data[np.newaxis, :, :]
    # res = backtesting(data, "MACD_RSI_OBV_Strategy")
    res = backtesting(data, "CSM_Strategy")

    # 清洗所有数据
    portfolio_cleaned = clean_data(res[0])  # 投资组合净值
    prices_cleaned = res[1]  # 价格数据
    actions_cleaned = res[2]  # 交易动作

    # 可视化
    show.plot_portfolio_value(portfolio_cleaned)
    show.plot_price_and_actions(prices_cleaned, actions_cleaned)
    print(show.calc_portfolio_metrics(portfolio_cleaned))
